﻿
/**
*ref:https://blog.csdn.net/summer_9527/article/details/96438935
*/

#include "ffmpeg_frame_any.h"
#ifdef __cplusplus
extern "C"
{
#endif
#include "libavformat/avformat.h"  
#include <libswscale/swscale.h>
#ifdef __cplusplus
}
#endif
static bool is_init = false;
static AVFormatContext *pFormatCtx;
static int videoStream = -1;
static struct SwsContext *img_convert_ctx;
static AVCodecContext *pCodecCtx;
static AVFrame *pFrameRGB;
static uint8_t *out_buffer;
static AVPacket *packet;
static AVFrame *pFrame;
static double r2d(AVRational r)
{
    return r.num == 0 || r.den == 0 ? 0. : (double)r.num / (double)r.den;
}

static int init(const char* inputFile)
{  
    AVCodec *pCodec;
    is_init = false;
    av_register_all();
    pFormatCtx = avformat_alloc_context();
    if (avformat_open_input(&pFormatCtx, inputFile, NULL, NULL) != 0) {
        return -1;
    }
    if (avformat_find_stream_info(pFormatCtx, NULL) < 0) {
        return -2;
    }
    
    for (auto i = 0; i < pFormatCtx->nb_streams; i++) {
        if (pFormatCtx->streams[i]->codec->codec_type == AVMEDIA_TYPE_VIDEO) {
            videoStream = i;
        }
    }

    if (videoStream == -1) {
        return -3;
    }

    pCodecCtx = pFormatCtx->streams[videoStream]->codec;
    pCodec = avcodec_find_decoder(pCodecCtx->codec_id);
    if (pCodec == NULL) {
        return -4;
    }
    if (avcodec_open2(pCodecCtx, pCodec, NULL) < 0) {
        return -5;
    }
    pFrame = av_frame_alloc();
    pFrameRGB = av_frame_alloc();
    img_convert_ctx = sws_getContext(pCodecCtx->width, pCodecCtx->height,
        pCodecCtx->pix_fmt, pCodecCtx->width, pCodecCtx->height,
        AV_PIX_FMT_RGB24, SWS_BICUBIC, NULL, NULL, NULL);

    int numBytes = avpicture_get_size(AV_PIX_FMT_RGB24, pCodecCtx->width, pCodecCtx->height);
    out_buffer = (uint8_t *)av_malloc(numBytes * sizeof(uint8_t));
    avpicture_fill((AVPicture *)pFrameRGB, out_buffer, AV_PIX_FMT_RGB24,
        pCodecCtx->width, pCodecCtx->height);

    int y_size = pCodecCtx->width * pCodecCtx->height;

    packet = (AVPacket *)malloc(sizeof(AVPacket)); //分配一个packet
    av_new_packet(packet, y_size); //分配packet的数据
    is_init = true;
    return 0;
}

static void deInit() 
{
    if (!is_init)
        return;
    is_init = false;
    av_free(out_buffer);
    av_free(pFrameRGB);
    avcodec_close(pCodecCtx);
    avformat_close_input(&pFormatCtx);
}

int ffmpeg_frame(const char* inputFile, int frameIndex, int(*func)(int height, int width, char* data))
{  
 /*   cv::Mat img;*/
    int ret = 0, got_picture;
    if (!is_init)
        ret = init(inputFile);
    if (ret < 0) 
    {
        return -1;
    }
    double sec = (double)frameIndex / r2d(pFormatCtx->streams[videoStream]->r_frame_rate);
    int64_t time_stamp = pFormatCtx->streams[videoStream]->start_time;
    double  time_base = r2d(pFormatCtx->streams[videoStream]->time_base);
    time_stamp += (int64_t)(sec / time_base + 0.5);
    double time_gap = (int64_t)((1.0 / r2d(pFormatCtx->streams[videoStream]->r_frame_rate) / time_base + 0.5)) * 0.5;
    // av_seek_frame方法，最后一个参数为AVSEEK_FLAG_BACKWARD时，会自动寻找前面的关键帧，一般会出现定位不准的情况
    ret = av_seek_frame(pFormatCtx, videoStream, time_stamp, AVSEEK_FLAG_BACKWARD);
    if (ret < 0)
        return -2;
    while (av_read_frame(pFormatCtx, packet) >= 0)
    {
        if (packet->stream_index == videoStream) {
            //判断是否为视频流，通过判断播放时间来确定是否为所需的帧
            if (abs(packet->pts - time_stamp) > time_gap)
            {
                ret = avcodec_decode_video2(pCodecCtx, pFrame, &got_picture, packet);
                av_free_packet(packet);
                continue;
            }
            ret = avcodec_decode_video2(pCodecCtx, pFrame, &got_picture, packet);
            if (ret < 0) {
                av_free_packet(packet);
                break;
            }

            if (got_picture) {
                sws_scale(img_convert_ctx,
                    (uint8_t const * const *)pFrame->data,
                    pFrame->linesize, 0, pCodecCtx->height, pFrameRGB->data,
                    pFrameRGB->linesize);
                pFrameRGB->height = 1080;
                pFrameRGB->width = 1920;
                func(pFrameRGB->height, pFrameRGB->width, (char*)pFrameRGB->data[0]);
                av_free_packet(packet);
                break;
            }
            av_free_packet(packet);
            break;
        }
    }
    return 0;
}